# 20250711
'''
给定 xxx_balanced
    xxx_localimg
    xxx_localimg_mask
    xxx_localimg_cannny
    xxx_localimg_normal    
训练时 
    构建 [bala , bala-mask] w/ [mask , black]
    redux = [ siglip(localimg) , siglip(canny/normal) ]
'''
import os,pdb,shutil,json
from util_for_os import ose,osj

will_save = {'data':[]}
save_json_path = 'local_data_20250711.json'

types = ['sleeve','collar','pockets']
data_root = '/mnt/nas/shengjie2/datasets'
get_ori_data_dir = lambda t : osj( data_root , f'cloth_{t}_balanced' )
get_localimg_data_dir = lambda t : osj( data_root , f'cloth_{t}_localimg' )
get_localimg_mask_data_dir = lambda t : osj( data_root , f'cloth_{t}_localimg_mask' )
get_localimg_canny_data_dir = lambda t : osj( data_root , f'cloth_{t}_localimg_canny' )
# get_localimg_normal_data_dir = lambda t : osj( data_root , f'cloth_{t}_localimg_normal' )
for choose_type in types:
    ori_data_dir = get_ori_data_dir(choose_type)
    localimg_data_dir = get_localimg_data_dir(choose_type)
    localimg_mask_data_dir = get_localimg_mask_data_dir(choose_type)
    localimg_canny_data_dir = get_localimg_canny_data_dir(choose_type)
    # localimg_normal_data_dir = get_localimg_normal_data_dir(choose_type)

    count = 0
    for entry in os.scandir(ori_data_dir):
        if not entry.is_file():continue
        filename = entry.name
        if not filename.endswith('.jpg'):continue

        count+=1
        print('\rprocess idx: ',count,end='',flush=True)

        ori_filepath = osj( ori_data_dir , filename )
        localimg_filepath = osj( localimg_data_dir , filename )
        localimg_mask_filepath = osj( localimg_mask_data_dir , filename.replace('.jpg','.png') )
        localimg_canny_filepath = osj( localimg_canny_data_dir , filename )
        # localimg_normal_filepath = osj( localimg_normal_data_dir , filename )

        assert ose(ori_filepath),ori_filepath
        assert ose(localimg_filepath),localimg_filepath
        assert ose(localimg_mask_filepath),localimg_mask_filepath
        assert ose(localimg_canny_filepath),localimg_canny_filepath
        # assert ose(localimg_normal_filepath),localimg_normal_filepath

        will_save['data'].append(
            {
                'ori_filepath':ori_filepath,
                'localimg_filepath':localimg_filepath,
                'localimg_mask_filepath':localimg_mask_filepath,
                'localimg_canny_filepath':localimg_canny_filepath,
                # 'localimg_normal_filepath':localimg_normal_filepath,
            }
        )
    print(f'process num({choose_type}) : ',count)


with open(save_json_path , 'w',encoding='utf-8') as f:
    json.dump(will_save , f  ,indent=4 , ensure_ascii=False)